package grpc

import (
	"context"
	"fmt"
	"runtime"
	"time"

	"gitee.com/zawei/yggdrasil/pkg/log"
	"gitee.com/zawei/yggdrasil/pkg/utils/xstrings"
	"google.golang.org/grpc"
)

// StreamServerInterceptor returns a new streaming grpcServer interceptor that adds zap.Logger to the context.
func logStreamServerInterceptor(slowThreshold time.Duration) grpc.StreamServerInterceptor {
	return func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) (err error) {
		var (
			startTime = time.Now()
			fields    = make([]interface{}, 0, 12)
			event     = "normal"
		)
		defer func() {
			cost := time.Since(startTime)
			if slowThreshold <= cost {
				event = "slow"
			}

			if rec := recover(); rec != nil {
				switch rec := rec.(type) {
				case error:
					err = rec
				default:
					err = fmt.Errorf("%v", rec)
				}
				stack := make([]byte, 4096)
				stack = stack[:runtime.Stack(stack, true)]
				fields = append(fields, "stack", xstrings.Bytes2str(stack))
				event = "recover"
			}

			fields = append(fields, "type", "stream", "method", info.FullMethod, "cost", cost, "event", event)

			if err != nil {
				fields = append(fields, "error", err.Error())
				log.Errorw("access", fields...)
				return
			}
			log.Infow("access", fields...)
		}()
		return handler(srv, stream)
	}
}

// UnaryServerInterceptor returns a new unary grpcServer interceptors that adds zap.Logger to the context.
func logUnaryServerInterceptor(slowThreshold time.Duration) grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
		var (
			startTime = time.Now()
			fields    = make([]interface{}, 0, 12)
			event     = "normal"
		)
		defer func() {
			cost := time.Since(startTime)
			if slowThreshold <= cost {
				event = "slow"
			}

			if rec := recover(); rec != nil {
				switch rec := rec.(type) {
				case error:
					err = rec
				default:
					err = fmt.Errorf("%v", rec)
				}
				stack := make([]byte, 4096)
				stack = stack[:runtime.Stack(stack, true)]
				fields = append(fields, "stack", xstrings.Bytes2str(stack))
				event = "recover"
			}
			fields = append(fields, "type", "unary", "method", info.FullMethod, "cost", cost, "event", event)

			if err != nil {
				fields = append(fields, "error", err)
				log.Errorw("access", fields...)
				return
			}
			log.Infow("access", fields...)
		}()
		return handler(ctx, req)
	}
}
